"""
通用WebSocket工具 - 基于Linus设计哲学
"WebSocket应该像普通路由一样简单" - 消除所有样板代码
"""

import asyncio
import time
from typing import Dict, Any, Optional, Callable
from gua import create_websocket_manager, ws_text, ws_json
from .logger import get_default_logger


class WebSocketHandler:
    """Linus式WebSocket处理器 - 消除所有重复的WebSocket样板代码"""
    
    def __init__(self, manager=None):
        # 使用传入的管理器或创建新的（向后兼容）
        self.ws_manager = manager or create_websocket_manager()
    
    async def handle_chat(
        self,
        scope: Dict[str, Any],
        receive: Callable,
        send: Callable,
        welcome_message: str = "欢迎连接到WebSocket聊天室",
        on_message: Optional[Callable] = None,
        manager=None
    ):
        # 使用传入的管理器或实例的管理器
        ws_manager = manager or self.ws_manager
        """处理聊天WebSocket - 统一所有聊天室逻辑"""
        
        client_ip = scope.get("client", [""])[0]
        
        try:
            await ws_text(send, "accept")
            ws_manager.add(send)
            
            # 发送欢迎消息
            await ws_json(send, {
                "type": "welcome",
                "message": welcome_message,
                "client_ip": client_ip,
                "timestamp": time.time()
            })
            
            # 广播新用户加入
            await ws_manager.broadcast({
                "type": "user_joined",
                "message": f"用户 {client_ip} 加入了聊天室",
                "timestamp": time.time()
            })
            
            while True:
                message = await receive()
                if message["type"] == "websocket.receive":
                    text = message.get("text", "")
                    if text.strip():
                        # 自定义消息处理或默认广播
                        if on_message:
                            await on_message(text, client_ip, send, ws_manager)
                        else:
                            await ws_manager.broadcast({
                                "type": "message",
                                "sender": client_ip,
                                "data": text,
                                "timestamp": time.time()
                            })
                elif message["type"] == "websocket.disconnect":
                    break
                    
        except Exception as e:
            logger = get_default_logger()
            logger.error(
                "WebSocket聊天处理器发生错误",
                event="websocket_chat_handler_error",
                error=str(e),
                error_type=type(e).__name__,
                client_ip=client_ip,
                user_impact="用户的聊天连接将被关闭"
            )
            try:
                await ws_json(send, {
                    "type": "error",
                    "message": "连接发生错误",
                    "timestamp": time.time()
                })
            except:
                pass
        finally:
            ws_manager.remove(send)
            # 广播用户离开
            try:
                await ws_manager.broadcast({
                    "type": "user_left",
                    "message": f"用户 {client_ip} 离开了聊天室",
                    "timestamp": time.time()
                })
            except:
                pass
    
    async def handle_upload_progress(
        self,
        scope: Dict[str, Any],
        receive: Callable,
        send: Callable,
        progress_callback: Optional[Callable] = None
    ):
        """处理文件上传进度WebSocket - 统一上传进度逻辑"""
        
        await ws_text(send, "accept")
        self.ws_manager.add(send)
        
        try:
            while True:
                message = await receive()
                if message["type"] == "websocket.receive":
                    command = message.get("text", "")
                    
                    if command.startswith("upload:"):
                        filename = command[7:]
                        
                        if progress_callback:
                            await progress_callback(filename, send)
                        else:
                            # 默认上传进度模拟
                            for progress in range(0, 101, 10):
                                await ws_json(send, {
                                    "type": "upload_progress",
                                    "filename": filename,
                                    "progress": progress,
                                    "timestamp": time.time()
                                })
                                await asyncio.sleep(0.2)
                            
                            await ws_json(send, {
                                "type": "upload_complete",
                                "filename": filename,
                                "message": "文件上传完成",
                                "timestamp": time.time()
                            })
        except Exception:
            pass
        finally:
            self.ws_manager.remove(send)
    
    async def handle_ai_chat(
        self,
        scope: Dict[str, Any],
        receive: Callable,
        send: Callable,
        ai_response_generator: Callable
    ):
        """处理AI对话WebSocket - 统一AI对话逻辑"""
        
        await ws_text(send, "accept")
        self.ws_manager.add(send)
        
        # 对话上下文
        conversation_context = []
        
        try:
            while True:
                message = await receive()
                if message["type"] == "websocket.receive":
                    user_message = message.get("text", "")
                    
                    # 添加到对话上下文
                    conversation_context.append({
                        "role": "user",
                        "message": user_message,
                        "timestamp": time.time()
                    })
                    
                    # 发送思考状态
                    await ws_json(send, {
                        "type": "thinking",
                        "status": "processing",
                        "message": "AI正在思考..."
                    })
                    
                    # 生成AI响应
                    ai_response = ai_response_generator(user_message, conversation_context)
                    
                    # 添加AI响应到上下文
                    conversation_context.append({
                        "role": "ai",
                        "message": ai_response,
                        "timestamp": time.time()
                    })
                    
                    # 流式发送AI响应
                    await ws_json(send, {
                        "type": "ai_response_start",
                        "user_message": user_message,
                        "total_length": len(ai_response),
                        "timestamp": time.time()
                    })
                    
                    # 逐字流式发送
                    current_response = ""
                    for i, char in enumerate(ai_response):
                        current_response += char
                        
                        # 简化的延迟计算
                        delay = 0.05 + (0.1 if char in '。！？；：' else 0)
                        if i > 0 and i % 10 == 0:
                            delay += 0.2
                        
                        await asyncio.sleep(delay)
                        
                        await ws_json(send, {
                            "type": "ai_response_chunk",
                            "content": char,
                            "index": i,
                            "total": len(ai_response),
                            "current_response": current_response,
                            "timestamp": time.time()
                        })
                    
                    # 发送完成状态
                    await ws_json(send, {
                        "type": "ai_response_complete",
                        "full_response": ai_response,
                        "total_chars": len(ai_response),
                        "timestamp": time.time()
                    })
                    
                    # 保持上下文在合理范围内
                    if len(conversation_context) > 10:
                        conversation_context = conversation_context[-8:]
                        
        except Exception:
            pass
        finally:
            self.ws_manager.remove(send)


class SimpleWebSocketHandler:
    """极简WebSocket处理器 - Linus式最简单实现"""
    
    @staticmethod
    async def echo_handler(scope, receive, send):
        """简单的回显处理器"""
        from gua import ws_text
        await ws_text(send, "accept")
        
        try:
            while True:
                message = await receive()
                if message["type"] == "websocket.receive":
                    text = message.get("text", "")
                    await ws_text(send, "send", f"回显: {text}")
                elif message["type"] == "websocket.disconnect":
                    break
        except Exception:
            pass
    
    @staticmethod
    async def time_handler(scope, receive, send):
        """时间推送处理器"""
        from gua import ws_text
        await ws_text(send, "accept")
        
        try:
            while True:
                await ws_text(send, "send", {
                    "type": "time",
                    "timestamp": time.time(),
                    "formatted": time.strftime("%Y-%m-%d %H:%M:%S")
                })
                await asyncio.sleep(1)
        except Exception:
            pass


# 全局实例 - 使用全局WebSocket管理器
from gua import ws_manager
ws_handler = WebSocketHandler(ws_manager)


def create_chat_websocket(*args, **kwargs):
    """便捷函数：创建聊天WebSocket"""
    return ws_handler.handle_chat(*args, **kwargs)


def create_upload_progress_websocket(*args, **kwargs):
    """便捷函数：创建上传进度WebSocket"""
    return ws_handler.handle_upload_progress(*args, **kwargs)


def create_ai_chat_websocket(*args, **kwargs):
    """便捷函数：创建AI对话WebSocket"""
    return ws_handler.handle_ai_chat(*args, **kwargs)